﻿/* Copyright (C) 2011 by John Cronin
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:

 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.

 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

using System;
using System.Collections.Generic;
using System.Text;

namespace tysos
{
    class ElfReader
    {
        struct Elf64_Ehdr
        {
            public UInt64 e_ident_1;
            public UInt64 e_ident_2;
            public UInt32 e_type_machine;
            public UInt32 e_machine;
            public UInt64 e_entry;
            public UInt64 e_phoff;
            public UInt64 e_shoff;
            public UInt32 e_flags;
            public UInt32 e_ehsize_phentsize;
            public UInt32 e_phnum_shentsize;
            public UInt32 e_shnum_shstrndx;
        }

        struct Elf64_Shdr
        {
            public UInt32 sh_name;
            public UInt32 sh_type;
            public UInt64 sh_flags;
            public UInt64 sh_addr;
            public UInt64 sh_offset;
            public UInt64 sh_size;
            public UInt32 sh_link;
            public UInt32 sh_info;
            public UInt64 sh_addralign;
            public UInt64 sh_entsize;
        }

        struct Elf64_Sym
        {
            public UInt32 st_name;
            public UInt32 st_info_other_shndx;
            public UInt64 st_value;
            public UInt64 st_size;
        }

        struct Elf64_Phdr
        {
            public UInt32 p_type;
            public UInt32 p_flags;
            public UInt64 p_offset;
            public UInt64 p_vaddr;
            public UInt64 p_paddr;
            public UInt64 p_filesz;
            public UInt64 p_memsz;
            public UInt64 p_align;
        }

        struct Elf64_Dyn
        {
            public Int64 d_tag;
            public UInt64 d_val;

            public const Int64 DT_NULL = 0;
            public const Int64 DT_PLTRELSZ = 2;
            public const Int64 DT_STRTAB = 5;
            public const Int64 DT_SYMTAB = 6;
            public const Int64 DT_RELA = 7;
            public const Int64 DT_RELASZ = 8;
            public const Int64 DT_RELAENT = 9;
            public const Int64 DT_SYMENT = 11;
            public const Int64 DT_PLTREL = 20;
            public const Int64 DT_JMPREL = 23;
        }

        struct Elf64_Rela
        {
            public UInt64 r_offset;
            public UInt64 r_info;
            public Int64 r_addend;

            public const UInt32 R_X86_64_64 = 1;
            public const UInt32 R_X86_64_GLOB_DAT = 6;
            public const UInt32 R_X86_64_JUMP_SLOT = 7;
        }

        private unsafe static Elf64_Ehdr* VerifyElf(ulong binary)
        {
            Elf64_Ehdr* ehdr = (Elf64_Ehdr*)binary;

            /* magic should equal 0x7f, 'E', 'L', 'F'
             * which is equal to 0x464c457f */
            uint magic = (uint)(ehdr->e_ident_1 & 0xffffffff);
            if (magic != 0x464c457f)
                throw new Exception("Not a valid Elf magic number: " + magic.ToString());

            /* EI_CLASS should equal ELFCLASS64 (= 2) */
            ulong ei_class = (ehdr->e_ident_1 >> 32) & 0xff;
            if (ei_class != 2)
                throw new Exception("Invalid EI_CLASS: " + ei_class.ToString());

            return ehdr;
        }

        public static ulong GetEntryPoint(ulong binary, ulong load_address)
        {
            unsafe
            {
                Elf64_Ehdr* ehdr = VerifyElf(binary);
                return ehdr->e_entry + load_address;
            }
        }

        public static ulong LoadModule(Virtual_Regions vreg, VirtMem vmem, SymbolTable stab, ulong binary, ulong binary_paddr, string name)
        {
            unsafe
            {
                Elf64_Ehdr* ehdr = VerifyElf(binary);

                /* Allocate a physical region to host the relocated binary */
                ulong loaded_size = GetLoadedSize(binary);
                ulong pie_base = vreg.Alloc(loaded_size, 0x1000, name + "_exec");

                /* Load up the sections marked PT_LOAD, store the PT_DYNAMIC section
                 * for later */
                ulong e_phentsize = (ulong)((ehdr->e_ehsize_phentsize >> 16) & 0xffff);
                uint e_phnum = (uint)(ehdr->e_phnum_shentsize & 0xffff);
                ulong prog_header = binary + ehdr->e_phoff;
                Elf64_Phdr* dynamic_phdr = null;

                for (uint i = 0; i < e_phnum; i++)
                {
                    Elf64_Phdr* cur_phdr = (Elf64_Phdr*)prog_header;

                    if (cur_phdr->p_type == 0x1)
                    {
                        /* We only load types of PT_LOAD (=1) */

                        /* All modules in tysos (as pointed to by the 'binary' heading are mapped directly
                         * from the physical location they have been loaded to by GRUB)
                         * This implies that the offset of any part of a file within a virtual page is the
                         * same as the offset within the physical page it is mapped from.
                         * 
                         * We can take advantage of this.
                         * If the offset of a section within the binary & 0xfff is equal to its load address & 0xfff
                         * then we can just map the physical addresses contained within the section to its load address,
                         * it they are different we need to map in some new pages at the load address and then copy
                         * the data into it
                         */

                        if (((cur_phdr->p_offset + binary_paddr) & 0xfff) == (cur_phdr->p_vaddr & 0xfff))
                        {
                            if(cur_phdr->p_filesz != cur_phdr->p_memsz)
                                throw new Exception("p_filesz != p_memsz - this is currently unsupported");

                            ulong page_offset = cur_phdr->p_vaddr & 0xfff;

                            ulong paddr_start = cur_phdr->p_offset + binary_paddr - page_offset;
                            ulong vaddr_start = pie_base + cur_phdr->p_vaddr - page_offset;
                            ulong length = util.align(cur_phdr->p_filesz + page_offset, 0x1000);

                            for (ulong j = 0; j < length; j += 0x1000)
                                vmem.map_page(vaddr_start + j, paddr_start + j);
                        }
                        else
                            throw new Exception("Binary not paged aligned - this is currently unsupported");
                    }
                    else if (cur_phdr->p_type == 0x2)
                    {
                        /* Store the PT_DYNAMIC section for later */
                        dynamic_phdr = cur_phdr;
                    }

                    prog_header += e_phentsize;
                }

                if (dynamic_phdr == null)
                    throw new Exception("Dynamic section not found");

                /* Now iterate through the dynamic table to identify the virtual addresses of the
                 * dynamic string and symbol tables */
                ulong dyn_sym_vaddr = 0;
                ulong dyn_str_vaddr = 0;
                ulong rela_tab_vaddr = 0;
                ulong rela_length = 0;
                ulong rela_entsize = 0;
                ulong sym_entsize = 0;
                ulong pltrela_tab_vaddr = 0;
                ulong pltrela_length = 0;
                ulong dtab_addr = binary + dynamic_phdr->p_offset;
                ulong dtab_entsize = 16;
                ulong dtab_end = dtab_addr + dynamic_phdr->p_filesz;

                while (dtab_addr < dtab_end)
                {
                    Elf64_Dyn* dyntab_entry = (Elf64_Dyn*)dtab_addr;

                    if (dyntab_entry->d_tag == Elf64_Dyn.DT_STRTAB)
                        dyn_str_vaddr = pie_base + dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_SYMTAB)
                        dyn_sym_vaddr = pie_base + dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_RELA)
                        rela_tab_vaddr = pie_base + dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_RELAENT)
                        rela_entsize = dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_RELASZ)
                        rela_length = dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_SYMENT)
                        sym_entsize = dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_PLTREL)
                    {
                        if (dyntab_entry->d_val != Elf64_Dyn.DT_RELA)
                            throw new Exception("PLT relocation section does not use rela relocations - currently not supported");
                    }
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_JMPREL)
                        pltrela_tab_vaddr = pie_base + dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_PLTRELSZ)
                        pltrela_length = dyntab_entry->d_val;
                    else if (dyntab_entry->d_tag == Elf64_Dyn.DT_NULL)
                        dtab_addr = dtab_end;

                    dtab_addr += dtab_entsize;
                }
                if (dyn_str_vaddr == 0)
                    throw new Exception("Dynamic string table not found");
                if (dyn_sym_vaddr == 0)
                    throw new Exception("Dynamic symbol table not found");
                if (rela_tab_vaddr == 0)
                    throw new Exception("Rela table not found");
                if (rela_length == 0)
                    throw new Exception("Rela table length not found");
                if (rela_entsize == 0)
                    throw new Exception("Rela entsize not found");
                if (sym_entsize == 0)
                    throw new Exception("Symbol entsize not found");

                /* Dump the dynamic sections */
                Formatter.WriteLine("Dynamic section entries:", Program.arch.DebugOutput);
                Formatter.Write(" load address: ", Program.arch.DebugOutput);
                Formatter.Write(pie_base, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);              
                Formatter.Write(" DT_DYNSTR:    ", Program.arch.DebugOutput);
                Formatter.Write(dyn_str_vaddr, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_DYNSYM:    ", Program.arch.DebugOutput);
                Formatter.Write(dyn_sym_vaddr, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_RELA:      ", Program.arch.DebugOutput);
                Formatter.Write(rela_tab_vaddr, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_RELAENT:   ", Program.arch.DebugOutput);
                Formatter.Write(rela_entsize, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_RELASZ:    ", Program.arch.DebugOutput);
                Formatter.Write(rela_length, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_SYMENT:    ", Program.arch.DebugOutput);
                Formatter.Write(sym_entsize, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_JMPREL:    ", Program.arch.DebugOutput);
                Formatter.Write(pltrela_tab_vaddr, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.Write(" DT_PLTRELSZ:  ", Program.arch.DebugOutput);
                Formatter.Write(pltrela_length, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);

                /* Now iterate through the dynamic relocations and fix up what
                 * needs to be fixed up
                 */

                ulong cur_rela_vaddr = rela_tab_vaddr;
                ulong rela_tab_end = rela_tab_vaddr + rela_length;
                while (cur_rela_vaddr < rela_tab_end)
                {
                    Elf64_Rela* cur_rela = (Elf64_Rela*)cur_rela_vaddr;

                    uint r_sym = (uint)((cur_rela->r_info >> 32) & 0xffffffff);
                    uint r_type = (uint)(cur_rela->r_info & 0xffffffff);

                    ulong* r_loc = (ulong*)(pie_base + cur_rela->r_offset);

                    if (r_type == Elf64_Rela.R_X86_64_64)
                    {
                        /* R_X86_64_64 relocations are 64 bit and point to S + A */
                        ulong sym_addr = GetSymbolAddr(dyn_sym_vaddr, dyn_str_vaddr, sym_entsize, pie_base, stab, r_sym);
                        *r_loc = util.Add(sym_addr, cur_rela->r_addend);

                        unchecked
                        {
                            Formatter.Write("Relocating object at ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)r_loc, "X", Program.arch.DebugOutput);
                            Formatter.Write(" type ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)r_type, "d", Program.arch.DebugOutput);
                            Formatter.Write(" S: ", Program.arch.DebugOutput);
                            Formatter.Write(sym_addr, "X", Program.arch.DebugOutput);
                            Formatter.Write(" A: ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)cur_rela->r_addend, "X", Program.arch.DebugOutput);
                            Formatter.Write(" to S + A: ", Program.arch.DebugOutput);
                            Formatter.Write(util.Add(sym_addr, cur_rela->r_addend), "X", Program.arch.DebugOutput);
                            Formatter.WriteLine(Program.arch.DebugOutput);
                        }
                    }
                    else if (r_type == Elf64_Rela.R_X86_64_GLOB_DAT)
                    {
                        /* R_X86_64_GLOB_DAT relocations are 64 bit and point to S */
                        ulong sym_addr = GetSymbolAddr(dyn_sym_vaddr, dyn_str_vaddr, sym_entsize, pie_base, stab, r_sym);
                        *r_loc = sym_addr;

                        unchecked
                        {
                            Formatter.Write("Relocating object at ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)r_loc, "X", Program.arch.DebugOutput);
                            Formatter.Write(" type ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)r_type, "d", Program.arch.DebugOutput);
                            Formatter.Write(" S: ", Program.arch.DebugOutput);
                            Formatter.Write(sym_addr, "X", Program.arch.DebugOutput);
                            Formatter.Write(" to S: ", Program.arch.DebugOutput);
                            Formatter.Write(sym_addr, "X", Program.arch.DebugOutput);
                            Formatter.WriteLine(Program.arch.DebugOutput);
                        }
                    }
                    else if (r_type == Elf64_Rela.R_X86_64_JUMP_SLOT)
                    {
                        /* R_X86_64_JUMP_SLOT relocations are 64 bit and point to S */
                        ulong sym_addr = GetSymbolAddr(dyn_sym_vaddr, dyn_str_vaddr, sym_entsize, pie_base, stab, r_sym);
                        *r_loc = sym_addr;

                        unchecked
                        {
                            Formatter.Write("Relocating object at ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)r_loc, "X", Program.arch.DebugOutput);
                            Formatter.Write(" type ", Program.arch.DebugOutput);
                            Formatter.Write((ulong)r_type, "d", Program.arch.DebugOutput);
                            Formatter.Write(" S: ", Program.arch.DebugOutput);
                            Formatter.Write(sym_addr, "X", Program.arch.DebugOutput);
                            Formatter.Write(" to S: ", Program.arch.DebugOutput);
                            Formatter.Write(sym_addr, "X", Program.arch.DebugOutput);
                            Formatter.WriteLine(Program.arch.DebugOutput);
                        }
                    }
                    else
                        throw new Exception("Relocation type not supported (" + r_type.ToString() + ")");

                    cur_rela_vaddr += rela_entsize;
                }

                /* Iterate through the pltrela sections, if provided */
                if (pltrela_tab_vaddr != 0)
                {
                    if (pltrela_length == 0)
                        throw new Exception("DT_PLTRELSZ not found");

                    ulong cur_pltrela_vaddr = pltrela_tab_vaddr;
                    ulong pltrela_tab_end = pltrela_tab_vaddr + pltrela_length;

                    while (cur_pltrela_vaddr < pltrela_tab_end)
                    {
                        Elf64_Rela* cur_rela = (Elf64_Rela*)cur_pltrela_vaddr;

                        uint r_sym = (uint)((cur_rela->r_info >> 32) & 0xffffffff);
                        uint r_type = (uint)(cur_rela->r_info & 0xffffffff);

                        ulong* r_loc = (ulong*)(pie_base + cur_rela->r_offset);

                        if (r_type == Elf64_Rela.R_X86_64_JUMP_SLOT)
                        {
                            /* R_X86_64_JUMP_SLOT relocations are 64 bit and point to S */
                            ulong sym_addr = GetSymbolAddr(dyn_sym_vaddr, dyn_str_vaddr, sym_entsize, pie_base, stab, r_sym);
                            *r_loc = sym_addr;
                        }
                        else
                            throw new Exception("Relocation type not supported (" + r_type.ToString() + ")");

                        cur_pltrela_vaddr += rela_entsize;
                    }
                }

                return pie_base;
            }
        }

        private unsafe static ulong GetSymbolAddr(ulong dyn_sym_vaddr, ulong dyn_str_vaddr, ulong sym_entsize, ulong pie_base, SymbolTable stab, uint r_sym)
        {
            ulong sym_addr = dyn_sym_vaddr + (ulong)r_sym * sym_entsize;
            Elf64_Sym* sym = (Elf64_Sym*)sym_addr;

            uint st_shndx = (uint)(sym->st_info_other_shndx >> 16);

            Formatter.Write("Resolving symbol: ", Program.arch.DebugOutput);
            Formatter.Write(new string((sbyte*)(dyn_str_vaddr + sym->st_name)), Program.arch.DebugOutput);

            if (st_shndx == 0)
            {
                /* SHN_UNDEF identifies an undefined symbol - we need to fetch it from the global symbol table */
                string sym_name = new string((sbyte*)(dyn_str_vaddr + sym->st_name));
                ulong addr = stab.GetAddress(sym_name);
                Formatter.Write(" found in symbol table, address: ", Program.arch.DebugOutput);
                Formatter.Write(addr, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                return addr;
            }
            else
            {
                /* Its a symbol defined within this file */
                Formatter.Write(" locally defined as: ", Program.arch.DebugOutput);
                Formatter.Write(pie_base + sym->st_value, "X", Program.arch.DebugOutput);
                Formatter.WriteLine(Program.arch.DebugOutput);
                return pie_base + sym->st_value;
            }
        }

        private unsafe static ulong GetLoadedSize(ulong binary)
        {
            Elf64_Ehdr* ehdr = (Elf64_Ehdr*)binary;

            ulong e_phentsize = (ulong)((ehdr->e_ehsize_phentsize >> 16) & 0xffff);
            uint e_phnum = (uint)(ehdr->e_phnum_shentsize & 0xffff);

            ulong prog_header = binary + ehdr->e_phoff;
            ulong loaded_size = 0;
            for (uint i = 0; i < e_phnum; i++)
            {
                Elf64_Phdr* phdr = (Elf64_Phdr*)prog_header;

                /* We only load types of PT_LOAD (=1) */
                if (phdr->p_type == 0x1)
                {
                    ulong cur_end = phdr->p_vaddr + phdr->p_memsz;

                    if (cur_end > loaded_size)
                        loaded_size = cur_end;
                }

                prog_header += e_phentsize;
            }

            return loaded_size;
        }

        public static void LoadSymbols(SymbolTable stab, ulong binary, ulong symbol_adjust)
        {
            unsafe
            {
                Elf64_Ehdr* ehdr = VerifyElf(binary);

                /* Iterate through the section headers and find the symbol table */
                ulong sect_header = binary + ehdr->e_shoff;
                bool found = false;
                ulong sym_tab_start = 0;
                ulong str_tab_sect_header = 0;
                ulong sym_tab_entsize = 0;
                ulong sym_tab_size = 0;
                
                for (uint i = 0; i < (ehdr->e_shnum_shstrndx & 0xffff); i++)
                {
                    Elf64_Shdr* shdr = (Elf64_Shdr*)sect_header;

                    /* SHT_SYMTAB is 2 */
                    if (shdr->sh_type == 2)
                    {
                        sym_tab_start = binary + shdr->sh_offset;
                        str_tab_sect_header = binary + ehdr->e_shoff + shdr->sh_link * (ulong)((ehdr->e_phnum_shentsize >> 16) & 0xffff);
                        sym_tab_entsize = shdr->sh_entsize;
                        sym_tab_size = shdr->sh_size;
                        found = true;
                        break;
                    }

                    sect_header += (ulong)((ehdr->e_phnum_shentsize >> 16) & 0xffff);
                }
                if (!found)
                    throw new Exception("Symbol table not found");

                /* Load the string table */
                Elf64_Shdr* str_tab = (Elf64_Shdr*)str_tab_sect_header;

                /* Iterate through all symbols, loading them into the symbol table */
                ulong sym_tab_offset = 0;
                while (sym_tab_offset < sym_tab_size)
                {
                    Elf64_Sym* cur_sym = (Elf64_Sym*)(sym_tab_start + sym_tab_offset);

                    ulong name_addr = binary + str_tab->sh_offset + (ulong)cur_sym->st_name;
                    ulong st_info = (ulong)(cur_sym->st_info_other_shndx & 0xff);
                    ulong st_shndx = (ulong)(cur_sym->st_info_other_shndx >> 16);

                    /* we only load symbols with STB_GLOBAL (=1) binding and
                     * of type STT_OBJECT (=1) or STT_FUNC (=2)
                     * 
                     * In st_info, binding is the high 4 bits, symbol type is low 4
                     * 
                     * Therefore we are looking for 00010001b or 00010010b
                     * which is 0x11 or 0x12
                     */

                    if ((st_info != 0x11) && (st_info != 0x12))
                    {
                        sym_tab_offset += sym_tab_entsize;
                        continue;
                    }

                    /* We do not want symbols with st_shndx == SHN_UNDEF (=0)
                     * We ignore symbols with st_shndx == SHN_COMMON (=0xfff2) 
                     */

                    if ((st_shndx == 0x0) || (st_shndx == 0xfff2))
                    {
                        sym_tab_offset += sym_tab_entsize;
                        continue;
                    }
                   
                    string sym_name = new string((sbyte*)name_addr);
                    ulong sym_addr = cur_sym->st_value + symbol_adjust;

                    stab.Add(sym_name, sym_addr);

                    sym_tab_offset += sym_tab_entsize;
                }
            }

            Formatter.WriteLine("ELF file parsed successfully", Program.arch.DebugOutput);
        }
    }
}
